import SocketServer
from socket import *
from struct import unpack,pack
import pcap
import threading
import time
from datetime import datetime, timedelta

from NSGBOT_libISP import getCountryFromIP

MAX_TAG_LEN = 30

FLOW_DIC = {} # (srcip,dstip,dstport) : ((ftime,ltime),[taglist])
BL_LIST = {}  # DsrIP:DsrPort  dictionary 
IRC_LIST = {} # Identified IRC Server List (by USER, PRIVMSG tags)
DB_BLACKLIST = {}  # DB (ip:name)
DB_WHITELIST = {}  # DB (ip:name)

# Global settings
DETECT_THRESHOLD = 0.7
MIN_CMD = 30
TIMEOUT = 3600
MAX_PAYLOAD_SIZE = 1500
SIZE_OF_INFO_T = 32
SIZE_OF_INFO2_T = 64
DELIM = "|"
MAX_PERIOD = 360                # FLOW TIME_TO_LIVE (second)
THRESHOLD = 0.5
THRESHOLD_SUM = 0.9

# TAG COM
TAG_USER = 11
TAG_MODE = 12
TAG_NICK = 13
TAG_USERHOST = 14
TAG_JOIN = 15
TAG_PING = 21
TAG_PONG = 22
TAG_PRIVMSG = 31
TAG_NOTICE = 32
TAG_NOTICE_AUTH = 33

# Result BL_LIST index: [country, bottype, score_max, score_cur, taglist, ftime, ltime]
RESULT_COUNTRY = 0
RESULT_BOTTYPE = 1
RESULT_SCORE_MAX = 2
RESULT_SCORE_CUR = 3
RESULT_TAGLIST = 4
RESULT_FTIME   = 5
RESULT_LTIME   = 6

# Bottype
BOTTYPE_IRC = 0
BOTTYPE_WEB = 1
BOTTYPE_P2P = 3

# Score function (index, value)
SCORE_MORE = 10
SCORE_BLACKLIST = 0
SCORE_FOREIGNIP = 1
SCORE_PING = 2
SCORE_PONG = 3
SCORE_PINGPONG = 4
SCORE_DECREASE = 99

SCORE = {
    SCORE_BLACKLIST : 0.5,
    SCORE_FOREIGNIP : 0.2,
    SCORE_PING      : 0.5,
    SCORE_PONG      : 0.6,
    SCORE_PINGPONG  : 0.7,
    SCORE_DECREASE  : -0.005,
}

 
def has_member(list, mem):
    for index in list:
        if (mem == index): return True
    return False

def show_list(LIST):
    for key in LIST.keys():
        print key, LIST[key]

class BotAnalyzer(threading.Thread):
    def run(self):
        print "Starting IRC bot analyzer..."
        
        while 1:
            print "\n===================== LOG ==================="
            keys = FLOW_DIC.keys()
            now = time.time()
            for key in keys:
                ((ftime,ltime),tag_list) = FLOW_DIC[key]
                (init, conn, chat) = self.tag_count(tag_list)
                total = len(tag_list)
                
                # Time check (delete dead flow)
                # count is effective, when bigger than MAX_TAG_LEN
                (sip, dip, dport) = key
                if self.check_blacklist(dip) == True:
                    self.add_score(key, SCORE_BLACKLIST)
                    self.update_ltime(key, ltime)

                if total < MIN_CMD:
                    # recent tags
                    last_update = float(ltime)
                    if (now - last_update) < MAX_PERIOD:
                        continue

                    # very old flow
                    # Check forgeign ip
                    if self.check_foreignIP(dip) == True:
                        self.add_score(key, SCORE_FOREIGNIP)
                        self.update_ltime(key, ltime)

                    # delete key in flow_dic
                    del FLOW_DIC[key]
                    continue

                print key, "%d %d %d / %d" % (init, conn, chat, total)
                # calculate 
                init_rate = float(init)/total
                conn_rate = float(conn)/total
                
                if (init_rate > DETECT_THRESHOLD):
                    dtype = self.detection_type(tag_list, 10) 

                elif (conn_rate > DETECT_THRESHOLD):
                    score_index = self.detection_type(tag_list, 20)
                    self.add_score(key, score_index)
                    self.update_ltime(key, ltime)
                    del FLOW_DIC[key]
                else:
                    continue

            # Display BL_LIST
            print "SrcIP\t\tDstIP\t\tDstPort\tCountry\tBType\tScore(Max)\tScore(Cur)\tTAG List\tFirst Time\tLast Time"
            for key in BL_LIST.keys():
                (sip,dip,dport) = key
                (country,btype,score_max,score_cur,taglist,ftime,ltime) = BL_LIST[key]
                print "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s" % (sip, dip, dport, country, btype, score_max, score_cur, taglist,ftime,ltime)
                self.dec_score(key)

            time.sleep(30)

    def add_tag(self, key, score_index):
        # add tag into tag list
        data = BL_LIST[key]
        score = data[RESULT_SCORE_CUR]
        taglist = data[RESULT_TAGLIST]

        if has_member(taglist, score_index) == False:
            if score_index != SCORE_DECREASE:
                taglist.append(score_index)
            data[RESULT_SCORE_CUR] = data[RESULT_SCORE_CUR] + SCORE[score_index]
        else:
        # already exist in taglist
            additional = SCORE[score_index] / SCORE_MORE
            data[RESULT_SCORE_CUR] = data[RESULT_SCORE_CUR] + additional

        #check max score
        if data[RESULT_SCORE_MAX] < data[RESULT_SCORE_CUR]:
            data[RESULT_SCORE_MAX] = data[RESULT_SCORE_CUR]

        # del current score < 0
        if data[RESULT_SCORE_CUR] < 0:
            self.del_score(key)

    def add_score(self, key, score_index):
        # add score into BL_LIST[key]
        (sip, dip, dport) = key
        if BL_LIST.has_key(key) == False:
            ((ftime,ltime),tag_list) = FLOW_DIC[key]
            ################
            # Add new info
            # Country Btype Score(Max) Score(Cur) TagList Ftime Ltime
            BL_LIST[key] = [getCountryFromIP(dip), [], 0, 0, [], ftime, ltime] 

        # update tag list
        self.add_tag(key, score_index)

        # update Bottype
        self.update_bottype(key)

    def del_score(self, key):
        del BL_LIST[key]
        # To Do
        # insert into LOG_DB

    def dec_score(self, key):
        self.add_tag(key, SCORE_DECREASE)

    def update_ltime(self, key, ltime):
        data = BL_LIST[key]
        data[RESULT_LTIME] = ltime

    def update_bottype(self, key):
        (sip, dip, dport) = key
        # IRC Bot
        if IRC_LIST.has_key(dip) == True:
            print "bottype dip:", dip
            data = BL_LIST[key]
            if has_member(data[RESULT_BOTTYPE], BOTTYPE_IRC) == False:
                data[RESULT_BOTTYPE].append(BOTTYPE_IRC)
        # Web Bot
        # P2P Bot
                

    def check_foreignIP(self, dip):
        country = getCountryFromIP(dip)
        if country != "KR":
            return True    # suspicious
        return False       # no suspicious

    def check_blacklist(self, dip):
        if DB_BLACKLIST.has_key(dip) == True:
            return True
        return False


    def has_irc_cmd(self, tag_list):
        # login, privmsg is real irc command
        # check tag_list has login or privmsg
        for tag in tag_list:
            if tag == 11:        #TAG_USER
                return 1
            elif tag == 31:      #TAG_PRIVMSG
                return 1
        return 0

    def update_bottype(self, dip):
        if IRC_LIST.has_key(dip) == True:
            return 1
        else:
            return 0

    def tag_count(self, tag_list):
        init = 0
        conn = 0
        chat = 0
        #
        #while(len(tag_list) > MIN_CMD):
        #    tag_list.pop(0)
            
        for tag in tag_list:
            index = tag / 10
            if index == 1:
                init = init + 1
            elif index == 2:
                conn = conn + 1
            elif index == 3:
                chat = chat + 1
        return (init, conn, chat)

    def tag_mean(self, tag_list, tag_type):
        sum = 0
        count = 0
        for tag in tag_list:
            if (tag / tag_type) == 1:
                sum = sum + tag
                count = count + 1
        print "sum:%d, count:%d" % (sum,count)
        return float(sum)/count

    def detection_type(self, tag_list, tag_type):
        if tag_type == 10: # init
            pass
        elif tag_type == 20: # conn
            mean = self.tag_mean(tag_list, tag_type)
            if (mean == 21):
                return SCORE_PING  # by PING
            elif (mean == 22):
                return SCORE_PONG  # by PONG
            else:
                return SCORE_PINGPONG  # by PING-PONG

    def reduce_tag_list(self, key):
        (time, tag_list) = FLOW_DIC[key]
        for index in range(MIN_CMD/2):
            tag_list.pop(0)


class ThreadedMyUDPHandler(SocketServer.BaseRequestHandler):
    """
    This class works similar to the TCP handler class, except that
    self.request consists of a pair of data and client socket, and since
    there is no connection the client address must be given explicitly
    when sending data back via sendto().
    """

    def handle(self):
        data =  self.request[0].strip()
        socket = self.request[1]
        cur_thread = threading.currentThread()
        
        # RECV processing
        count = len(data) / SIZE_OF_INFO_T
        print "Data length:(%d),(num of %d)" % (len(data), count)

        pos = 0
        info_t = ""
        for index in range(count):
            pos = index * SIZE_OF_INFO_T
            #print "Current pos(%d),index(%d)" %(pos, index)

            if pos > len(data):
                print "Error pos:", pos
                break

            info_t = data[pos:pos+SIZE_OF_INFO_T]
            
            #info_t parsing
            timet = unpack('Q',info_t[0:8])[0]
            saddr = pcap.ntoa(unpack('i',info_t[8:12])[0])
            daddr = pcap.ntoa(unpack('i',info_t[12:16])[0])



            # check white list
            if DB_WHITELIST.has_key(daddr):
                continue

            dport = unpack('H',info_t[16:18])[0]
            tag = ord(info_t[18])
            payload = info_t[18:32]

            # TEST
            if daddr == "82.94.222.186":
                print "#detect 82.94.222.186"
                print timet, tag
                print payload
            # IRC JOIN TAG (add to IRC_LIST)
            if tag == TAG_JOIN:
                if IRC_LIST.has_key(daddr) == True:
                    #IRC_LIST[daddr] = payload
                    show_list(IRC_LIST)
                else:
                    print "No IRC Server LIST"
                    print payload

            key = (saddr, daddr, dport)

            # already exist
            if FLOW_DIC.has_key(key) == True:
                # update FLOW_DIC
                ((ftime,ltime),tag_list) = FLOW_DIC[key]
                tag_list.append(tag)
                FLOW_DIC[key] = ((ftime,timet),tag_list)
            # Fist time    
            else:
                FLOW_DIC[key] = ((timet,timet),[tag])
                
            # IRC_LIST update
            if tag == TAG_USER or tag == TAG_PRIVMSG or tag == TAG_NOTICE_AUTH:
                if IRC_LIST.has_key(daddr) == False:
                    IRC_LIST[daddr] = [dport]
                else:
                    dport_list = IRC_LIST[daddr]
                    if has_member(dport_list, dport) == False:
                        dport_list.append(dport)

                    print "##### IRC Server List ####"
                    for dip in IRC_LIST.keys():
                        print dip, IRC_LIST[dip]

            #print saddr, daddr, dport, tag, payload

class ThreadedUDPServer(SocketServer.ThreadingMixIn, SocketServer.UDPServer):
    pass


# Discovery Nation, ISP
class DiscoverISP(threading.Thread):
    def run(self):
        print "Starting DiscoverISP..."

def initDB():
    # init DB file
    # db.blacklist
    # db.whitelist
    fp = open("db.blacklist")
    for index in fp:
        (ip,name) = index.split("|")
        DB_BLACKLIST[ip] = name
    fp.close()

    fp = open("db.whitelist")
    for index in fp:
        (ip,name) = index.split("|")
        DB_WHITELIST[ip] = name
    fp.close()

if __name__ == "__main__":
    #init
    initDB()

    # start BotAnalyzer thread
    ba_thread = BotAnalyzer()
    ba_thread.start()

    HOST, PORT = "59.27.92.27", 9999
    BL_PORT = 9999

    server = ThreadedUDPServer((HOST, PORT), ThreadedMyUDPHandler)

    ip, port = server.server_address

    server_thread = threading.Thread(target=server.serve_forever)
    server_thread.setDaemon(False)
    server_thread.start()

    print "Server loop running in thread:", server_thread.getName()

